# !/usr/bin/env python
# -*- coding:utf-8 -*-
'''简单3d浏览

历史:
    2018-11-30 姚彧 移植
'''

__all__ = ['ySimple3dWidget']

from PyQt5 import QtWidgets
import vtk
from vtk.qt.QVTKRenderWindowInteractor import QVTKRenderWindowInteractor
import numpy as np
from helper.vector_dimension_convertor import to3d

global_matrix = np.array([(1, 0, 0),
                          (0, 0, -1),
                          (0, 1, 0)])

class ComponentData(object):
    def __init__(self, attribs):
        assert isinstance(attribs, dict)
        for key in attribs:
            setattr(self, key, attribs[key])

def mkVtkIdList(it):
    vil = vtk.vtkIdList()
    for i in it:
        vil.InsertNextId(int(i))
    return vil

def create_mesh(component: ComponentData):
    points, faces = component.data
    points = to3d(points)
    points = points.dot(global_matrix)
    middle_point = points.mean(0)

    vtkpoints = vtk.vtkPoints()
    for pt in points:
        vtkpoints.InsertNextPoint(*pt)

    cells = vtk.vtkCellArray()
    for face in faces:
        cells.InsertNextCell(mkVtkIdList(face))

    scalars = vtk.vtkFloatArray()
    for i, _ in enumerate(points):
        scalars.InsertTuple1(i, i)

    # 将点和线加入到数据集
    polydata = vtk.vtkPolyData()
    polydata.SetPoints(vtkpoints)
    polydata.SetPolys(cells)
    #
    polydata.GetPointData().SetScalars(scalars)

    # mapper
    mapper = vtk.vtkPolyDataMapper()
    mapper.SetInputData(polydata)
    mapper.SetScalarRange(polydata.GetScalarRange())

    # actor
    actor = vtk.vtkActor()
    actor.SetMapper(mapper)
    actor.GetProperty().SetRepresentationToWireframe()
    return actor, middle_point

class ySimple3dWidget(object):
    def __init__(self, container_widget):
        self.actor = None
        self.gridlayout = QtWidgets.QGridLayout(container_widget)
        self.vtkWidget = QVTKRenderWindowInteractor(container_widget)
        self.gridlayout.addWidget(self.vtkWidget, 0, 0, 1, 1)

        self.renderer = vtk.vtkRenderer()
        self.renderer.SetBackground( 0.1, 0.2, 0.4 )
        self.vtkWidget.GetRenderWindow().AddRenderer(self.renderer)
        self.iren = self.vtkWidget.GetRenderWindow().GetInteractor()

        # 自定义交互模式
        style = vtk.vtkInteractorStyleTrackballCamera()
        style.SetDefaultRenderer(self.renderer)
        self.iren.SetInteractorStyle(style)

    def add_component(self, attribs):
        self.renderer.Clear()
        if self.actor:
            self.renderer.RemoveActor(self.actor)
        actor, middle_point = create_mesh(ComponentData(attribs))
        self.actor = actor
        self.renderer.AddActor(actor)
        camera = self.renderer.GetActiveCamera()
        focal_point = camera.GetFocalPoint()
        position_point = camera.GetPosition()
        position_point = middle_point + np.array(position_point) - focal_point
        camera.SetPosition(*position_point)
        camera.SetFocalPoint(*middle_point)
        self.vtkWidget.GetRenderWindow().Render()
